Attention Mechanism Comparison¶
In [2]:
import torch
import numpy as np
import matplotlib.pyplot as plt
import seaborn as sns
import plotly.graph_objects as go
import plotly.express as px
from plotly.subplots import make_subplots
import plotly.io as pio
import pandas as pd
from datetime import datetime
import warnings
warnings.filterwarnings('ignore')
from transformers import (
AutoTokenizer, AutoModel,
pipeline, AutoConfig
)
from datasets import load_dataset
from bertviz import model_view, head_view, neuron_view
from bertviz.transformers_neuron_view import BertModel
plt.style.use('seaborn-v0_8')
sns.set_palette("husl")
torch.manual_seed(42)
np.random.seed(42)
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
print(f"Using device: {device}")
Using device: cpu
In [3]:
import html
import re
def load_sentences(n_samples=100, min_words=10, max_words=20):
"""Load sentences from AG News dataset"""
dataset = load_dataset("ag_news", split="train")
label_names = {0: "World", 1: "Sports", 2: "Business", 3: "Sci/Tech"}
sentences = []
domains = []
samples_per_category = n_samples // 4
for label in range(4):
filtered = dataset.filter(lambda x: x['label'] == label)
shuffled = filtered.shuffle(seed=42)
count = 0
for item in shuffled:
if count >= samples_per_category:
break
# decode HTML and clean text
text = html.unescape(item['text'])
text = re.sub(r'<[^>]+>', '', text) # Remove HTML tags
text = re.sub(r'\s+', ' ', text).strip() # Normalize whitespace
# Check length
word_count = len(text.split())
if min_words <= word_count <= max_words and len(text) > 20:
sentences.append(text)
domains.append(label_names[label])
count += 1
print(f"\n Loaded {len(sentences)} sentences from AG News:")
for label_name in label_names.values():
count = domains.count(label_name)
print(f" {label_name}: {count} sentences")
print(f"\n Sample sentences:")
for i in range(min(3, len(sentences))):
print(f" [{domains[i]}] {sentences[i][:80]}...")
return sentences, domains
TEST_SENTENCES, SENTENCE_DOMAINS = load_sentences(n_samples=100, min_words=10, max_words=20)
Loaded 100 sentences from AG News: World: 25 sentences Sports: 25 sentences Business: 25 sentences Sci/Tech: 25 sentences Sample sentences: [World] Somalis vie to be new president Twenty-eight candidates are approved to contest ... [World] Agency pleads for hostage release Care International appeals on Arabic televisio... [World] Clinton recovering after heart op Former US President Bill Clinton's heart bypas...
In [4]:
MODELS = {
"BERT": {
"name": "bert-base-uncased",
"description": "Original BERT",
"color": "#FF6B6B"
},
"DistilBERT": {
"name": "distilbert-base-uncased",
"description": "Compressed BERT",
"color": "#4ECDC4"
},
"RoBERTa": {
"name": "roberta-base",
"description": "Optimized BERT",
"color": "#45B7D1"
}
}
In [5]:
def load_model_and_tokenizer(model_name):
"""
Load model and tokenizer for a given model name.
"""
tokenizer = AutoTokenizer.from_pretrained(model_name)
model = AutoModel.from_pretrained(
model_name,
output_attentions=True,
return_dict=True).to(device)
config = AutoConfig.from_pretrained(model_name)
print(f"Modèle {model_name} loaded")
return model, tokenizer,config
models_data = {}
for model_name, model_info in MODELS.items():
model, tokenizer, config = load_model_and_tokenizer(model_info["name"])
models_data[model_name] = {
"model": model,
"tokenizer": tokenizer,
"config": config,
"description": model_info["description"],
"color": model_info["color"]
}
print("✅ All models loaded")
Modèle bert-base-uncased loaded Modèle distilbert-base-uncased loaded
Some weights of RobertaModel were not initialized from the model checkpoint at roberta-base and are newly initialized: ['pooler.dense.bias', 'pooler.dense.weight'] You should probably TRAIN this model on a down-stream task to be able to use it for predictions and inference.
Modèle roberta-base loaded ✅ All models loaded
In [6]:
def analyze_model_architecture(models_data):
"""Compare architecture of different models."""
arch_data = []
for model_name, data in models_data.items():
config = data["config"]
model = data["model"]
total_params = sum(p.numel() for p in model.parameters())
trainable_params = sum(p.numel() for p in model.parameters() if p.requires_grad)
arch_info = {
"Model": model_name,
"Layers": config.num_hidden_layers,
"Hidden Size": config.hidden_size,
"Attention Heads": config.num_attention_heads,
"Vocab Size": config.vocab_size,
"Total Params": f"{total_params:,}",
"Trainable Params": f"{trainable_params:,}",
}
arch_data.append(arch_info)
print(f"\n{model_name}:")
print(f" Layers: {config.num_hidden_layers}")
print(f" Hidden Size: {config.hidden_size}")
print(f" Attention Heads: {config.num_attention_heads}")
print(f" Vocab Size: {config.vocab_size:,}")
print(f" Parameters: {total_params:,}")
return pd.DataFrame(arch_data)
architecture_df = analyze_model_architecture(models_data)
print("\n Comparison table:")
display(architecture_df)
BERT: Layers: 12 Hidden Size: 768 Attention Heads: 12 Vocab Size: 30,522 Parameters: 109,482,240 DistilBERT: Layers: 6 Hidden Size: 768 Attention Heads: 12 Vocab Size: 30,522 Parameters: 66,362,880 RoBERTa: Layers: 12 Hidden Size: 768 Attention Heads: 12 Vocab Size: 50,265 Parameters: 124,645,632 Comparison table:
| Model | Layers | Hidden Size | Attention Heads | Vocab Size | Total Params | Trainable Params | |
|---|---|---|---|---|---|---|---|
| 0 | BERT | 12 | 768 | 12 | 30522 | 109,482,240 | 109,482,240 |
| 1 | DistilBERT | 6 | 768 | 12 | 30522 | 66,362,880 | 66,362,880 |
| 2 | RoBERTa | 12 | 768 | 12 | 50265 | 124,645,632 | 124,645,632 |
In [7]:
def analyze_sentence_attention(sentence, model_name, models_data):
"""Analyze attention patterns for a given sentence and model."""
print(f"Analyze attention for: '{sentence}'")
print(f"Model: {model_name}")
tokenizer = models_data[model_name]["tokenizer"]
model = models_data[model_name]["model"]
inputs = tokenizer(sentence, return_tensors="pt", truncation=True)
tokens = tokenizer.convert_ids_to_tokens(inputs["input_ids"][0])
print(f"Tokens ({len(tokens)}): {tokens[:10]}{'...' if len(tokens) > 10 else ''}")
# Prediction with attention
with torch.no_grad():
outputs = model(**inputs)
attentions = outputs.attentions
num_layers = len(attentions)
num_heads = attentions[0].shape[1]
seq_length = attentions[0].shape[-1]
print(f"Attention shape: {num_layers} layers, {num_heads} heads, {seq_length} tokens")
return {
"tokens": tokens,
"attentions": attentions,
"inputs": inputs,
"num_layers": num_layers,
"num_heads": num_heads
}
test_sentence = TEST_SENTENCES[1]
print(f" Test on: '{test_sentence}'\n")
attention_results = {}
for model_name in models_data.keys():
attention_results[model_name] = analyze_sentence_attention(
test_sentence, model_name, models_data
)
print("-" * 50)
Test on: 'Agency pleads for hostage release Care International appeals on Arabic television for the release of its Iraq director, Margaret Hassan.' Analyze attention for: 'Agency pleads for hostage release Care International appeals on Arabic television for the release of its Iraq director, Margaret Hassan.' Model: BERT Tokens (25): ['[CLS]', 'agency', 'plead', '##s', 'for', 'hostage', 'release', 'care', 'international', 'appeals']... Attention shape: 12 layers, 12 heads, 25 tokens -------------------------------------------------- Analyze attention for: 'Agency pleads for hostage release Care International appeals on Arabic television for the release of its Iraq director, Margaret Hassan.' Model: DistilBERT Tokens (25): ['[CLS]', 'agency', 'plead', '##s', 'for', 'hostage', 'release', 'care', 'international', 'appeals']... Attention shape: 6 layers, 12 heads, 25 tokens -------------------------------------------------- Analyze attention for: 'Agency pleads for hostage release Care International appeals on Arabic television for the release of its Iraq director, Margaret Hassan.' Model: RoBERTa Tokens (26): ['<s>', 'A', 'gency', 'Ġple', 'ads', 'Ġfor', 'Ġhostage', 'Ġrelease', 'ĠCare', 'ĠInternational']... Attention shape: 12 layers, 12 heads, 26 tokens --------------------------------------------------
In [8]:
from IPython.display import display, HTML
def create_attention_visualizations(sentence, model_name, models_data):
"""Create attention visualizations with BertViz"""
tokenizer = models_data[model_name]["tokenizer"]
model = models_data[model_name]["model"]
inputs = tokenizer(sentence, return_tensors="pt", truncation=True, max_length=512)
tokens = tokenizer.convert_ids_to_tokens(inputs["input_ids"][0])
inputs = {k: v.to(device) for k, v in inputs.items()}
with torch.no_grad():
outputs = model(**inputs)
attentions = outputs.attentions
print(f"Tokens: {len(tokens)} tokens")
print(f"Attention: {len(attentions)} layers, {attentions[0].shape[1]} heads")
print("Model View (all attention heads):")
model_view(attentions, tokens)
print("Head View (per head detail):")
head_view(attentions, tokens)
return {
"tokens": tokens,
"attentions": attentions,
"num_layers": len(attentions),
"num_heads": attentions[0].shape[1]
}
visualization_results = {}
for model_name in models_data.keys():
print(f"\n{'='*60}")
print(f"VISUALIZING {model_name}")
print(f"{'='*60}")
result = create_attention_visualizations(test_sentence, model_name, models_data)
visualization_results[model_name] = result
============================================================ VISUALIZING BERT ============================================================ Tokens: 25 tokens Attention: 12 layers, 12 heads Model View (all attention heads):
Head View (per head detail):
Layer:
============================================================ VISUALIZING DistilBERT ============================================================ Tokens: 25 tokens Attention: 6 layers, 12 heads Model View (all attention heads):